#!/usr/bin/env python3
import os, sys
import numpy as np
import matplotlib.pylab as plt
from astropy.io import fits
from astropy.stats import SigmaClip

sigmaclip = SigmaClip(sigma=5, maxiters=5)

reg_names = ['A (DD)', 'B (TID)', 'C (TID)', 'D (DD)', 'E (NONE)', 'F (NONE)']
reg = [ [206,  4506, 182,  4382],       # A
        [4698, 8998, 168,  4368],       # B
        [4993, 7767, 5009, 9059],       # C
        [1444, 4344, 4859, 9073],       # D
        [8222, 9078, 4903, 9703],       # E
        [82,   1172, 4717, 9151] ]      # F


def getTScaleFactor(T_old, T_new):
    return (T_new/T_old)**3 * np.exp(6400.*(1/T_old - 1/T_new))

def dark_rescale(darke, T_old, T_new):
    scale_factor = getTScaleFactor(T_old, T_new)
    return darke * scale_factor

def print_darke_stats(darke):
    darke_clip = sigmaclip(darke)
    dc_mean = np.mean(darke_clip)
    # print('==> np.sum(darke>0.08) = {}'.format(np.sum(darke>0.08)))
    # print('==> np.sum(darke>0.16) = {}'.format(np.sum(darke>0.16)))
    idx = darke < 25    # remove hot pixels
    frac08 = 1.0*np.sum(darke[idx]>0.08) / darke.size
    frac16 = 1.0*np.sum(darke[idx]>0.16) / darke.size

    print('mean dc  = {} e-/pix/s'.format(dc_mean))
    print('frac08   = {}'.format(frac08))
    print('frac16   = {}'.format(frac16))

def print_regional_darke_stats(darke):
#    darke_clip = sigmaclip(darke)
#    dc_mean = np.mean(darke_clip)
#    idx = darke < 25    # remove hot pixels
#    frac08 = 1.0*np.sum(darke[idx]>0.08) / darke.size
#    frac16 = 1.0*np.sum(darke[idx]>0.16) / darke.size

    for i in range(len(reg)):
        py0, py1, px0, px1 = reg[i]
        regdarke = darke[px0:px1, py0:py1]
        darke_clip = sigmaclip(regdarke)
        dc_mean = np.mean(darke_clip)
        idx = regdarke < 25    # remove hot pixels
        frac08 = 1.0*np.sum(regdarke[idx]>0.08) / regdarke.size
        frac16 = 1.0*np.sum(regdarke[idx]>0.16) / regdarke.size
        print('## Region @{} darke statistical results:'.format(reg_names[i]))
        print('\tmean dc  = {:.4E} e-/pix/s'.format(dc_mean))
        print('\tfrac08   = {:.4E}'.format(frac08))
        print('\tfrac16   = {:.4E}'.format(frac16))

def plot_rescaled_dark(darke, T_old, T_new):
    print('--> plotting darke at {}K and rescaled to {}K'.format(T_old, T_new))
    darke_new = dark_rescale(darke, T_old, T_new)

    # idx_old = np.abs(darke) < 50
    # idx_new = np.abs(darke_new) < 50
    # print('**************************************')
    # print('before rescale: dc_mean  = {}'.format( np.mean(darke) ))
    # print('before rescale: dc_median= {}'.format( np.median(darke) ))
    # print('before rescale: frac08   = {}'.format( np.sum(darke[idx_old]>0.08) / np.sum(idx_old) ))
    # print('before rescale: frac16   = {}'.format( np.sum(darke[idx_old]>0.16) / np.sum(idx_old) ))
    # print('after rescale: dc_mean   = {}'.format( np.mean(darke_new)) )
    # print('after rescale: dc_median = {}'.format( np.median(darke_new)) )
    # print('after rescale: frac08    = {}'.format( np.sum(darke_new[idx_new]>0.08) / np.sum(idx_new) ))
    # print('after rescale: frac16    = {}'.format( np.sum(darke_new[idx_new]>0.16) / np.sum(idx_new) ))
    # print('**************************************')

    print('**************************************')
    print('@ before rescaled:')
#    print_darke_stats(darke)
    print_regional_darke_stats(darke)
    print('@ after rescaled:')
#    print_darke_stats(darke_new)
    print_regional_darke_stats(darke_new)
    print('**************************************')


#    plt.hist(darke[idx_old], bins=100, histtype='step', log=True, label='DC@'+str(T_old)+'K')
#    plt.hist(darke_new[idx_new], bins=100, histtype='step', log=True, label='DC@'+str(T_new)+'K')
#    plt.legend(loc='best')
#    ax = plt.gca()
#    ax.set_xlabel(r'$e^{-}/pix/sec$')
#    plt.tight_layout()
#    plt.show()

if __name__=='__main__':
    if len(sys.argv) != 4:
        print('usage: {} darke_XXXK.fits Temp1 Temp2'.format(sys.argv[0]))
        sys.exit(0)
    darke = fits.getdata(sys.argv[1]).astype(float)
    T_old = float(sys.argv[2])
    T_new = float(sys.argv[3])
    plot_rescaled_dark(darke, T_old, T_new)
